"""This module provides tools and shortcuts for working with python iterators.

Python iterators, in combination with generators and collections,
provide efficient generic tools for traversing different types of data structures.

- IteratorCounter: count number of items traversed by a iterator
- partition: separate the iterator stream into several iterator streams, each with the same given part_size
- partition_evenly: separate an iterator stream with a known size into a given `nparts` number of iterator streams,
   each stream with nearly equal size
"""


from typing import Iterator, Collection, Union


class IteratorCounter(object):
    def __init__(self, iterator):
        self.count = 0
        self.iterator = iterator

    def __iter__(self):
        return self

    def __next__(self):
        try:
            val = next(self.iterator)
        except StopIteration:
            raise
        else:
            self.count += 1
            return val


def partition(iterator: Union[Collection, Iterator], part_size: int):
    """Separate items in a given iterator to several partitions,
    each partition having the same size, while the last partition might have a smaller size.

    Parameters
    ----------
    iterator : an Iterator or a Collection
        Items to be separate into parts
    part_size : int
        Size of each part

    Yields
    -------
    List
    """
    part = []
    for item in iterator:
        part.append(item)
        if len(part) == part_size:
            yield part
            part = []
    else:
        if len(part) > 0:
            yield part


def partition_evenly(iterator, nparts, total_size):
    raise NotImplementedError("later")
